from matplotlib import pyplot

from fitter import Fitter, FitException
from bfitter import BinnedFitter

class FitterPolynomial(Fitter):
    def __init__(self, x, y, order, par_names=None, xerr=None, yerr=None, title='Polynomial Fit', xlabel='', ylabel=''):
        if order < 0:
            raise FitException("Negative order polynomials are not supported by this fitter.")
        
        if len(x) < (order+1) or len(y) < (order+1):
            raise FitException("Need at least {0} points to fit with a polynomial of degree {1}.".format(order+1, order))
        
        self.order = order
        par_names = ['$p_{0}$'.format(i) for i in xrange(self.order+1)]

        self.fcn_str = ' + '.join(['{' + str(i) + '}*x**' + str(i) for i in xrange(self.order+1)])
        
        self.fcn_latex = r'$p^{0}(x)$'.format(self.order)
        self._par_text_offset = 0.325
        
        super(FitterPolynomial, self).__init__(x, y, par_names=par_names, xerr=xerr, yerr=yerr, fcn_str=self.fcn_str, title=title, xlabel=xlabel, ylabel=ylabel)
        
    def _set_initpars(self, frange):
        return [1.0 for i in xrange(self.order+1)]


class BinnedFitterPolynomial(BinnedFitter):
    def __init__(self, data, order, par_names=None):
        if order < 0:
            raise ValueError("Negative order polynomials are not supported by this fitter.")

        self.order = order
        par_names = ['$p_{0}$'.format(i) for i in xrange(self.order+1)]


        self.fcn_str = ' + '.join(['{' + str(i) + '}*x**' + str(i) for i in xrange(self.order+1)])
        
        self.fcn_latex = '$p^{0}(x)$'.format(self.order)
        self._par_text_offset = 0.325

        super(BinnedFitterPolynomial, self).__init__(data, par_names = par_names, fcn_str=self.fcn_str)
        
    def _set_initpars(self, frange):
        return [1.0 for i in xrange(self.order+1)]
